#!/usr/bin/env python
"""This script will download album art for a music collection and put it in a specific folder"""

from __future__ import print_function

## Location of the music collection
MUSIC_LOCATION = "/home/user/Music"
## Location of the cover images
COVERS_LOCATION = "/home/user/.covers"
## Music formats to check
AUDIO_FORMATS = ("ogg", "flac", "mp3")
## Image files to accept as covers
IMAGE_FORMATS = ("png", "jpg", "bmp", "gif")
## Should it print everything out?
VERBOSE = True
## Disallowed characters in filenames
DISALLOWED_CHARS = "\\/:<>?*|"
import os
import sys
import urllib2
import glob

# Import tagging library
try:
    import tagpy
except ImportError:
    print("This script requires tagpy in order to work. On ubuntu, you can install it with sudo apt-get install python-tagpy.")
    sys.exit(1)

try:
    import pylast
except ImportError:
    print("This script requires python bindings for last.fm. Getting them now.")
    pylast_urlfile = urllib2.urlopen("http://pylast.googlecode.com/svn/trunk/pylast.py") # Get pylast.py from pylast's svn
    pylast_file = open(sys.path[0] + "/pylast.py",'wb')
    pylast_file.write(pylast_urlfile.read())
    pylast_file.close() # Save it to pylast.py in the current directory
    import pylast

def get_music_files(music_location):
    """Get a list of music files in the location"""
    # Get a list of audio files by walking through the folder tree
    music_files = []
    for root, sub_folders, files in os.walk(music_location, followlinks=True):
        for file in files:
            if file.split(".")[-1].lower() in AUDIO_FORMATS:
                music_files.append(os.path.join(root, file))
    # Remove files that can't be read by tagpy
    for file in music_files:
        # Try opening with tagpy
        try:
            tagpy.FileRef(file)
        # If it fails, then it is not a music file
        except ValueError:
            print("{0} can't be read".format(file))
            music_files.remove(file)

    return music_files

def get_album(file_path):
    """Returns the artist and album of a music file as a tuple"""
    music = tagpy.FileRef(file_path).tag()
    # Read the artist name
    try:
        artist = music.artist
        if artist ==  "":
            artist = None
    except AttributeError: # If the tag doesn't exist
        artist = None
    # Read the album name
    try:
        album = music.album
        if album == "":
            album = None
    except AttributeError: # If the tag doesn't exist
        album = None
    return artist, album

def sanitize(name):
    """Replaces disallowed characters with an underscore"""
    if name == None:
        name = "Unknown"
    for character in DISALLOWED_CHARS:
        name = name.replace(character,'_')
    # Replace " with '
    name = name.replace('"', "'")

    return name

def cover_exists(artist, album, location):
    """Check if a cover exists for a particular album"""
    return glob.glob(u"{0}/{1}-{2}.*".format(location, sanitize(artist), sanitize(album)).replace("[", "\\[").replace("]", "\\]"))

def list_covers_to_get(music_location, covers_location):
    """Get a list of covers to download"""
    covers_to_get = []
    # Loop through the audio files in the music folder
    for file in get_music_files(music_location):
        # Get the artist and album from the file
        album = get_album(file)
        # If the artist is known, album is not already in the list and it doesn't already exist in the cover folder
        if album[0] is not None and album not in covers_to_get and not cover_exists(album[0], album[1], location=covers_location):
            covers_to_get.append(album)

    return covers_to_get

def download_cover_art(artist, album, download_location, size=3, artist_image=False):
    """Use last.fm for album art fetching, returns path"""
    ## Last.FM API key
    api_key = '5dde241299a6b25d28769972e966471b'
    network = pylast.get_lastfm_network(api_key = api_key)
    # If the album is None, we can use the artist images instead
    if album and not artist_image:
        # Get the album object using album info we have
        album_object = network.get_album(artist, album)
        # Get the url of the cover art
        try:
            cover_url = album_object.get_cover_image(size=size)
        except pylast.WSError:
            cover_url = None
    else:
        # Get artist object
        artist_object = network.get_artist(artist)
        # Get the url of the artist photograph
        try:
            cover_url = artist_object.get_images()[0].sizes.largesquare
        except IndexError:
            cover_url = None
    # Get the cover image and put it in the covers directory, if it has been found
    if cover_url != None:
        cover_file = urllib2.urlopen(cover_url)
        output_path = os.path.join(download_location, u"{0}-{1}.{2}".format(sanitize(artist),
                                                                           sanitize(album),
                                                                           cover_url.split('.')[-1]))
        # If the download location doesn't exist, create it
        if not os.path.exists(download_location):
            os.makedirs(download_location)
        with open(output_path, 'wb') as cover_output:
            cover_output.write(cover_file.read())
    elif artist_image == False:
        # Otherwise try getting an artist image
        output_path = download_cover_art(artist, album, download_location, artist_image=True)
    else: # Both artist and album art failed to download so give up
        output_path = None

    return output_path

def get_cover(artist, album, covers_location):
    """Downloads a cover for a particular album if it doesn't already exist and then returns the path"""
    cover = cover_exists(artist, album, covers_location)
    # If the cover doesn't exist
    if cover:
        cover = cover[0]
    else:
        cover = download_cover_art(artist, album, covers_location)

    return cover

def get_covers(music_location, covers_location, size=3, verbose=False):
    """Download all needed covers"""
    # Albums that need covers
    albums = list_covers_to_get(music_location, covers_location)
    # Total number of albums to get
    total = len(albums)
    current = 0
    # Initiate a list of failed covers
    failed_covers = []
    for album in albums:
        current += 1
        # If the album cover couldn't be downloaded
        if download_cover_art(album[0], album[1], covers_location, size) is None:
            failed_covers.append(album)
            # Print error
            print(u"{0} - {1} could not be downloaded :(".format(album[0], album[1]), file=sys.stderr)
        # If successful
        else:
            if verbose:
                print(u"Got {0} - {1} [{2}/{3}]".format(album[0], album[1] if album[1] != None else "Unknown", current, total))

    # Total failed downloads
    if len(failed_covers) > 0:
        print("{0} album cover(s) could not be downloaded.".format(len(failed_covers)), file=sys.stderr)

if __name__ == "__main__":
    get_covers(MUSIC_LOCATION, COVERS_LOCATION, verbose=VERBOSE)

